{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Title: #All Nodes Distance K in Binary Tree"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Difficulty: #Medium"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Category Title: #Algorithms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tag Slug: #tree #depth-first-search #breadth-first-search #binary-tree"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Name Translated: #树 #深度优先搜索 #广度优先搜索 #二叉树"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solution Name: distanceK"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Title: #二叉树中所有距离为 K 的结点"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Content:\n",
    "<p>给定一个二叉树（具有根结点&nbsp;<code>root</code>），&nbsp;一个目标结点&nbsp;<code>target</code>&nbsp;，和一个整数值 <code>k</code> 。</p>\n",
    "\n",
    "<p>返回到目标结点 <code>target</code> 距离为 <code>k</code> 的所有结点的值的列表。 答案可以以 <strong>任何顺序</strong> 返回。</p>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<ol>\n",
    "</ol>\n",
    "\n",
    "<p><strong>示例 1：</strong></p>\n",
    "\n",
    "<p><img src=\"https://s3-lc-upload.s3.amazonaws.com/uploads/2018/06/28/sketch0.png\" style=\"height: 429px; width: 500px;\" /></p>\n",
    "\n",
    "<pre>\n",
    "<strong>输入：</strong>root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, k = 2\n",
    "<strong>输出：</strong>[7,4,1]\n",
    "<strong>解释：</strong>所求结点为与目标结点（值为 5）距离为 2 的结点，值分别为 7，4，以及 1\n",
    "</pre>\n",
    "\n",
    "<p><strong>示例 2:</strong></p>\n",
    "\n",
    "<pre>\n",
    "<strong>输入:</strong> root = [1], target = 1, k = 3\n",
    "<strong>输出:</strong> []\n",
    "</pre>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong>提示:</strong></p>\n",
    "\n",
    "<ul>\n",
    "\t<li>节点数在&nbsp;<code>[1, 500]</code>&nbsp;范围内</li>\n",
    "\t<li><code>0 &lt;= Node.val &lt;= 500</code></li>\n",
    "\t<li><code>Node.val</code>&nbsp;中所有值 <strong>不同</strong></li>\n",
    "\t<li>目标结点&nbsp;<code>target</code>&nbsp;是树上的结点。</li>\n",
    "\t<li><code>0 &lt;= k &lt;= 1000</code></li>\n",
    "</ul>\n",
    "\n",
    "<p>&nbsp;</p>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Description: [all-nodes-distance-k-in-binary-tree](https://leetcode.cn/problems/all-nodes-distance-k-in-binary-tree/description/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solutions: [all-nodes-distance-k-in-binary-tree](https://leetcode.cn/problems/all-nodes-distance-k-in-binary-tree/solutions/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_cases = ['[3,5,1,6,2,0,8,null,null,7,4]\\n5\\n2', '[1]\\n1\\n3']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, x):\n",
    "#         self.val = x\n",
    "#         self.left = None\n",
    "#         self.right = None\n",
    "\n",
    "class Solution:\n",
    "    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:\n",
    "        def dfs(node, parent = None):\n",
    "            if node:\n",
    "                node.parent = parent\n",
    "                dfs(node.left, node)\n",
    "                dfs(node.right, node)\n",
    "        dfs(root)\n",
    "\n",
    "        queue = collections.deque([(target, 0)])\n",
    "        seen = {target}\n",
    "        while queue:\n",
    "            if queue[0][1] == k:\n",
    "                return [node.val for node, d in queue]\n",
    "            node, d = queue.popleft()\n",
    "            for neighbor in (node.left, node.right, node.parent):\n",
    "                if neighbor and neighbor not in seen:\n",
    "                    seen.add(neighbor)\n",
    "                    queue.append((neighbor, d+1))\n",
    "\n",
    "        return []\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, x):\n",
    "#         self.val = x\n",
    "#         self.left = None\n",
    "#         self.right = None\n",
    "\n",
    "class Solution:\n",
    "    def distanceK(self, root, target, K):\n",
    "        \"\"\"\n",
    "        :type root: TreeNode\n",
    "        :type target: TreeNode\n",
    "        :type K: int\n",
    "        :rtype: List[int]\n",
    "        \"\"\"\n",
    "        def getNdepth(n, dep):\n",
    "            if not n:\n",
    "                return []\n",
    "            que = collections.deque([n])\n",
    "            while len(que) > 0 and dep > 0:\n",
    "                dep -= 1\n",
    "                for i in range(len(que)):\n",
    "                    cur = que.popleft()\n",
    "                    if cur.left:\n",
    "                        que.append(cur.left)\n",
    "                    if cur.right:\n",
    "                        que.append(cur.right)\n",
    "            return list(map(lambda x: x.val, que))\n",
    "        \n",
    "        def findNode(n):\n",
    "            if n:\n",
    "                if n == target:\n",
    "                    return [n]\n",
    "                lRes = findNode(n.left)\n",
    "                if lRes:\n",
    "                    lRes.append(n)\n",
    "                    return lRes\n",
    "                rRes = findNode(n.right)\n",
    "                if rRes:\n",
    "                    rRes.append(n)\n",
    "                    return rRes\n",
    "            return None\n",
    "        \n",
    "        path = findNode(root)\n",
    "        ans = getNdepth(target, K)\n",
    "        for i in range(1, min(len(path), K)):\n",
    "            ans.extend(getNdepth(path[i].right if path[i].left == path[i - 1]\n",
    "                                 else path[i].left, K - 1 - i))\n",
    "        if K != 0 and K < len(path):\n",
    "            ans.append(path[K].val)\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, x):\n",
    "#         self.val = x\n",
    "#         self.left = None\n",
    "#         self.right = None\n",
    "\n",
    "class Solution:\n",
    "    def distanceK(self,root, target, K):\n",
    "        \"\"\"\n",
    "        :type root: TreeNode\n",
    "        :type target: TreeNode\n",
    "        :type K: int\n",
    "        :rtype: List[int]\n",
    "        \"\"\"\n",
    "        # DFS\n",
    "        conn = collections.defaultdict(list)\n",
    "        def connect(parent, child):\n",
    "            if parent and child:\n",
    "                conn[parent.val].append(child.val)\n",
    "                conn[child.val].append(parent.val)\n",
    "            if child.left: connect(child, child.left)\n",
    "            if child.right: connect(child, child.right)\n",
    "        connect(None, root)\n",
    "        # BFS\n",
    "        que = collections.deque()\n",
    "        que.append(target.val)\n",
    "        visited = set([target.val])\n",
    "        for k in range(K):\n",
    "            size = len(que)\n",
    "            for i in range(size):\n",
    "                node = que.popleft()\n",
    "                for j in conn[node]:\n",
    "                    if j not in visited:\n",
    "                        que.append(j)\n",
    "                        visited.add(j)\n",
    "        return list(que)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distanceK(self, root, target, K):\n",
    "        node = [[-1, 0, 0] for i in range(520)]\n",
    "        self.deep = 0\n",
    "\n",
    "        def dfs(t, depth, f):\n",
    "            if not t:\n",
    "                return 0\n",
    "            node[t.val] = [f, depth, 0]\n",
    "            if t == target:\n",
    "                node[t.val][2] = 1\n",
    "                dfs(t.left, depth + 1, t.val)\n",
    "                dfs(t.right, depth + 1, t.val)\n",
    "                self.deep = depth\n",
    "                return 1\n",
    "            else:\n",
    "                x = dfs(t.left, depth + 1, t.val)\n",
    "                y = dfs(t.right, depth + 1, t.val)\n",
    "                if x or y:\n",
    "                    node[t.val][2] = 1\n",
    "                    return 1\n",
    "                else:\n",
    "                    return 0\n",
    "\n",
    "        dfs(root, 0, 511)\n",
    "        ans = []\n",
    "        for i in range(len(node)):\n",
    "            if node[i][0] >= 0:\n",
    "                y = i\n",
    "                while node[y][2] == 0:\n",
    "                    y = node[y][0]\n",
    "                if node[y][1] - self.deep + node[y][1] - node[i][1] == -K:\n",
    "                    ans.append(i)\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class NewNode:\n",
    "    def __init__(self, x):\n",
    "        self.val = x\n",
    "        self.parent = None\n",
    "        self.left = None\n",
    "        self.right = None\n",
    "        self.flag = False\n",
    "\n",
    "class Solution:\n",
    "    def distanceK(self, root, target, K):\n",
    "        \"\"\"\n",
    "        :type root: TreeNode\n",
    "        :type target: TreeNode\n",
    "        :type K: int\n",
    "        :rtype: List[int]\n",
    "        \"\"\"\n",
    "        tmp_node = target\n",
    "        target = tmp_node.val\n",
    "        new_root = NewNode(root.val)\n",
    "        self.dfs(root, new_root)\n",
    "        target_node = []\n",
    "        self.dfs2(new_root, target, target_node)\n",
    "        self.res = []\n",
    "        self.dfs3(target_node[0], K)\n",
    "        return self.res\n",
    "\n",
    "    def dfs(self, root, new_root):\n",
    "        if not root:\n",
    "            return\n",
    "        if root.right:\n",
    "            new_right = NewNode(root.right.val)\n",
    "            new_right.parent = new_root\n",
    "            new_root.right = new_right\n",
    "            self.dfs(root.right, new_right)\n",
    "        if root.left:\n",
    "            new_left = NewNode(root.left.val)\n",
    "            new_left.parent = new_root\n",
    "            new_root.left = new_left\n",
    "            self.dfs(root.left, new_left)\n",
    "\n",
    "    def dfs2(self, new_root, target, target_node):\n",
    "        if not new_root:\n",
    "            return\n",
    "        if new_root.val == target:\n",
    "            target_node.append(new_root)\n",
    "            return\n",
    "        self.dfs2(new_root.left, target, target_node)\n",
    "        self.dfs2(new_root.right, target, target_node)\n",
    "\n",
    "    def dfs3(self, new_node, k):\n",
    "        new_node.flag = True\n",
    "        if k == 0:\n",
    "            self.res.append(new_node.val)\n",
    "        else:\n",
    "            if new_node.parent and new_node.parent.flag is False:\n",
    "                self.dfs3(new_node.parent, k-1)\n",
    "            if new_node.left and new_node.left.flag is False:\n",
    "                self.dfs3(new_node.left, k-1)\n",
    "            if new_node.right and new_node.right.flag is False:\n",
    "                self.dfs3(new_node.right, k-1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, x):\n",
    "#         self.val = x\n",
    "#         self.left = None\n",
    "#         self.right = None\n",
    "\n",
    "class Solution:\n",
    "    def findParent(self,root):\n",
    "        if root.left:\n",
    "            root.left.parent = root\n",
    "            self.findParent(root.left)\n",
    "        if root.right:\n",
    "            root.right.parent = root\n",
    "            self.findParent(root.right)\n",
    "    def dfs(self,root,K,res,visited):\n",
    "        \"\"\"\n",
    "        深度遍历，将满足条件的节点值加入到res中，将访问过的节点值加入到visited中\n",
    "        \"\"\"\n",
    "        if root==None or root.val in visited:\n",
    "            # 如果结点为空或者已访问过，直接返回\n",
    "            return\n",
    "        visited.append(root.val)\n",
    "        if K==0:\n",
    "            # 如果K==0，说明当前结点就是要找的，不必再遍历子树\n",
    "            res.append(root.val)\n",
    "        else:\n",
    "            self.dfs(root.left,K-1,res,visited)\n",
    "            self.dfs(root.right,K-1,res,visited)\n",
    "            self.dfs(root.parent,K-1,res,visited)\n",
    "        \n",
    "            \n",
    "           \n",
    "    def distanceK(self, root, target, K):\n",
    "        \"\"\"\n",
    "        :type root: TreeNode\n",
    "        :type target: TreeNode\n",
    "        :type K: int\n",
    "        :rtype: List[int]\n",
    "        \"\"\"\n",
    "        root.parent = None\n",
    "        self.findParent(root)\n",
    "        res = []   # 存放符合条件节点的值\n",
    "        visited = []  # 存放已访问节点的值\n",
    "        self.dfs(target,K,res,visited);\n",
    "        \n",
    "        return res\n",
    "        \n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "class TreeNode:\n",
    "    def __init__(self, x):\n",
    "        self.val = x\n",
    "        self.left = None\n",
    "        self.right = None\n",
    "\n",
    "class Solution:\n",
    "\n",
    "    def distanceK(self, root: TreeNode, target: TreeNode, K: int) -> List[int]:\n",
    "        white, gray, black = -1, 0, 1\n",
    "\n",
    "        def dfs(node, parent):\n",
    "            if node is not None:\n",
    "                node.p = parent\n",
    "                node.color = white\n",
    "                dfs(node.left, node)\n",
    "                dfs(node.right, node)\n",
    "        dfs(root, None)\n",
    "\n",
    "        queue = []\n",
    "        queue.append((target, 0))\n",
    "        result = []\n",
    "\n",
    "        while queue != []:\n",
    "            node, distance = queue.pop(0)\n",
    "            if distance == K:\n",
    "                result.append(node.val)\n",
    "            adj = [node.left, node.right, node.p]\n",
    "            for v in adj:\n",
    "                if v is not None and v.color == white:\n",
    "                    v.color = gray\n",
    "                    queue.append((v, distance+1))\n",
    "            node.color = black\n",
    "        return result\n",
    "\n",
    "\n",
    "              \n",
    "        \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distanceK(self, root: TreeNode, target: TreeNode, K: int) -> List[int]:\n",
    "        res, self.lookup = [], dict()\n",
    "        self.__build_graph(None, root)\n",
    "        # 图擦创建之后,bfs搜索,需要seen集合记录访问过的,不然3-5再从5到3,步数都是2\n",
    "        step, seen = 0, set()\n",
    "        from collections import deque\n",
    "        queue = deque()\n",
    "        # 注意queue放的不再是root,而是从target节点开始走步\n",
    "        queue.append(target)\n",
    "        seen.add(target)\n",
    "        # 不光queue不为空,还需要没超过步数\n",
    "        while queue and step <= K:\n",
    "            size = len(queue)\n",
    "            for i in range(size):\n",
    "                node = queue.popleft()\n",
    "                # 如果当前已经是步数k了,就没必要继续遍历node的邻居了\n",
    "                if step == K: res.append(node.val)\n",
    "\n",
    "                # 拿出当前node所能走的下一步\n",
    "                neighbors = self.lookup.get(node, set())\n",
    "                for nei in neighbors:\n",
    "                    if nei in seen: continue\n",
    "                    queue.append(nei)\n",
    "                    seen.add(nei)\n",
    "            step += 1\n",
    "\n",
    "        return res\n",
    "\n",
    "    def __build_graph(self, parent, child):\n",
    "        # 构建无向图,parent-child的边和child-parent的边\n",
    "        if parent:\n",
    "            self.lookup.setdefault(parent, set())\n",
    "            self.lookup[parent].add(child)\n",
    "\n",
    "            self.lookup.setdefault(child, set())\n",
    "            self.lookup[child].add(parent)\n",
    "        # 根节点添加相邻节点后,遍历其左右节点\n",
    "        if child.left: self.__build_graph(child, child.left)\n",
    "        if child.right: self.__build_graph(child, child.right)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, x):\n",
    "#         self.val = x\n",
    "#         self.left = None\n",
    "#         self.right = None\n",
    "\n",
    "class Solution:\n",
    "    def distanceK(self, root: TreeNode, target: TreeNode, K: int) -> List[int]:\n",
    "\n",
    "        def build_graph(parent, child, neighbor):\n",
    "            if not child:\n",
    "                return\n",
    "            if parent:\n",
    "                neighbor[parent.val].append(child.val)\n",
    "                neighbor[child.val].append(parent.val)\n",
    "\n",
    "            build_graph(child, child.left, neighbor)\n",
    "            build_graph(child, child.right, neighbor)\n",
    "        from collections import defaultdict\n",
    "        nb = defaultdict(list)\n",
    "        build_graph(None, root, nb)\n",
    "        \n",
    "        res = [target.val]\n",
    "        visited = set(res)\n",
    "        for _ in range(K):\n",
    "            res = [n for t in res \\\n",
    "            for n in nb[t] \\\n",
    "            if n not in visited]\n",
    "\n",
    "            visited |= set(res) \n",
    "        return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, x):\n",
    "#         self.val = x\n",
    "#         self.left = None\n",
    "#         self.right = None\n",
    "\n",
    "class Solution:\n",
    "    def distanceK(self, root: TreeNode, target: TreeNode, K: int) -> List[int]:\n",
    "        self.res=[]\n",
    "        def depth(root,layer):\n",
    "            if root is None: return\n",
    "            if layer==K:self.res.append(root.val)\n",
    "            depth(root.left,layer+1) \n",
    "            depth(root.right,layer+1) \n",
    "        def find(root,path):\n",
    "            #if root is None: return\n",
    "            if root==target:return path\n",
    "            else:\n",
    "                if root.left:\n",
    "                    a=find(root.left,path+'0') \n",
    "                    if a:return a\n",
    "                if root.right:\n",
    "                    b=find(root.right,path+'1')\n",
    "                    if b:return b\n",
    "        def main(root,path,index):\n",
    "            if not root:return\n",
    "            if index==K:self.res.append(root.val)\n",
    "            if path:\n",
    "                if path[0]=='0':\n",
    "                    main(root.left,path[1:],index-1)\n",
    "                    main(root.right,'',index+1)\n",
    "                else: \n",
    "                    main(root.left,'',index+1)\n",
    "                    main(root.right,path[1:],index-1)\n",
    "            else:\n",
    "                main(root.left,'',index+1)\n",
    "                main(root.right,'',index+1) \n",
    "\n",
    "        if root==target:depth(root,0)\n",
    "        else:\n",
    "            path=find(root,'')\n",
    "            main(root,path,len(path))\n",
    "        return self.res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, x):\n",
    "#         self.val = x\n",
    "#         self.left = None\n",
    "#         self.right = None\n",
    "\n",
    "class Solution:\n",
    "    def distance(self,x0,y0,x1,y1):# 根据两个节点的坐标计算距离\n",
    "        res = 0\n",
    "        # 如果两个节点不在同一层\n",
    "        # 首先让layer较大的节点\n",
    "        # 即较低的节点上移，变成其父节点\n",
    "        # 对于任意节点，x,y\n",
    "        # 其父节点坐标为 x-1,y//2\n",
    "        # 相应地，距离+1\n",
    "        while x0 > x1:                                \n",
    "            x0 -= 1\n",
    "            y0 //= 2\n",
    "            res += 1\n",
    "        while x0 < x1:\n",
    "            x1 -= 1\n",
    "            y1 //= 2\n",
    "            res += 1\n",
    "        while y0 != y1:\n",
    "        # 当两节点处于同一层时\n",
    "        # 需要比较其pos\n",
    "        # 当pos相等时，两节点即重合\n",
    "        # 否则两个节点同步上移\n",
    "        # 由于其layer相等，只需考虑pos\n",
    "        # 对于任意 pos=y 的节点而言\n",
    "        # 其父节点pos = y//2\n",
    "        # 每次上移，距离+2\n",
    "            y0 //= 2\n",
    "            y1 //= 2\n",
    "            res += 2\n",
    "        return res\n",
    "\n",
    "    def distanceK(self, root: TreeNode, target: TreeNode, K: int) -> List[int]:\n",
    "        node_val = {}                                # 用于记录节点的坐标和值的对应关系\n",
    "        cur_level = [(0,0,root)]                         # 层序遍历的队列，（0,0）为根的坐标\n",
    "        x,y = 0,0                                    # 记录目标值所对应坐标\n",
    "        while cur_level:\n",
    "            layer,pos,node = cur_level.pop(0)            # 队列弹出的值分别为层数、位置和节点\n",
    "            node_val[(layer,pos)] = node.val # 记录节点的坐标和值的对应关系# 即二叉树的坐标化\n",
    "            if node == target:           # 找到并记录target对应的节点坐标\n",
    "                x = layer\n",
    "                y = pos\n",
    "            if node.left:   # 如果节点有左节点 # layer = layer + 1 # pos = pos * 2\n",
    "                cur_level.append((layer+1,pos*2,node.left))\n",
    "            if node.right: # 如果节点有右节点  # layer = layer + 1 # pos = pos * 2 + 1       \n",
    "                cur_level.append((layer+1,pos*2+1,node.right))\n",
    "        #利用distance函数 # 计算所有节点与target节点的距离 # 并返回距离为K的节点值\n",
    "        return [node_val[p] for p in node_val if self.distance(p[0],p[1],x,y)==K]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def distanceK(self, root: TreeNode, target: TreeNode, K: int) -> List[int]:\n",
    "        if K == 0:\n",
    "            return [target.val]\n",
    "\n",
    "        edge_dict = defaultdict(list)\n",
    "\n",
    "        tree_node_stack = list()\n",
    "        tree_node_stack.append(root)\n",
    "\n",
    "        while tree_node_stack:\n",
    "            cur_node = tree_node_stack.pop()\n",
    "            if cur_node.left:\n",
    "                tree_node_stack.append(cur_node.left)\n",
    "                edge_dict[cur_node.val].append(cur_node.left.val)\n",
    "                edge_dict[cur_node.left.val].append(cur_node.val)\n",
    "            if cur_node.right:\n",
    "                tree_node_stack.append(cur_node.right)\n",
    "                edge_dict[cur_node.val].append(cur_node.right.val)\n",
    "                edge_dict[cur_node.right.val].append(cur_node.val)\n",
    "\n",
    "        edge_node_stack = list()\n",
    "        edge_node_stack.append((target.val, K))\n",
    "        res_list = []\n",
    "        visited = set()\n",
    "        visited.add(target.val)\n",
    "\n",
    "        while edge_node_stack:\n",
    "            cur_edge_node = edge_node_stack.pop()\n",
    "            if cur_edge_node[1] == 1:\n",
    "                for value in edge_dict[cur_edge_node[0]]:\n",
    "                    if value not in visited:\n",
    "                        visited.add(value)\n",
    "                        res_list.append(value)\n",
    "            else:\n",
    "                for value in edge_dict[cur_edge_node[0]]:\n",
    "                    if value not in visited:\n",
    "                        visited.add(value)\n",
    "                        edge_node_stack.append((value, cur_edge_node[1]-1))\n",
    "\n",
    "        return res_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, x):\n",
    "#         self.val = x\n",
    "#         self.left = None\n",
    "#         self.right = None\n",
    "\n",
    "from collections import defaultdict\n",
    "\n",
    "class Solution:\n",
    "    def distanceK(self, root: TreeNode, target: TreeNode, K: int) -> List[int]:\n",
    "        dict_K = defaultdict(list)\n",
    "        queue = [root]\n",
    "        while queue:\n",
    "            node = queue.pop(0)\n",
    "            if node.left:\n",
    "                queue.append(node.left)\n",
    "                dict_K[node.val].append(node.left.val)\n",
    "                dict_K[node.left.val].append(node.val)\n",
    "\n",
    "            if node.right:\n",
    "                queue.append(node.right)\n",
    "                dict_K[node.val].append(node.right.val)\n",
    "                dict_K[node.right.val].append(node.val)\n",
    "        # print(dict_K)\n",
    "\n",
    "        bfs_queue = [target.val]\n",
    "        all_nodes = [target.val]\n",
    "        distance = 0\n",
    "\n",
    "        while bfs_queue:\n",
    "            if distance == K:\n",
    "                break\n",
    "                \n",
    "            length = len(bfs_queue)\n",
    "            distance += 1\n",
    "            for _ in range(length):\n",
    "                node = bfs_queue.pop(0)\n",
    "                next_list = dict_K.get(node, [])\n",
    "                for t in next_list:\n",
    "                    if t not in all_nodes:\n",
    "                        bfs_queue.append(t)\n",
    "                        all_nodes.append(t)\n",
    "            print(distance, bfs_queue)\n",
    "            \n",
    "\n",
    "        return bfs_queue\n",
    "\n",
    "\n",
    "\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, x):\n",
    "#         self.val = x\n",
    "#         self.left = None\n",
    "#         self.right = None\n",
    "\n",
    "class Solution:\n",
    "    def distanceK(self, root: TreeNode, target: TreeNode, K: int) -> List[int]:\n",
    "        import collections\n",
    "        graph = collections.defaultdict(list)\n",
    "        def build_graph(parent, child):\n",
    "            if parent:\n",
    "                graph[parent].append(child)\n",
    "                graph[child].append(parent)\n",
    "            if child.left:\n",
    "                build_graph(child, child.left)\n",
    "            if child.right:\n",
    "                build_graph(child, child.right)\n",
    "        \n",
    "        build_graph(None, root)\n",
    "        queue = collections.deque()\n",
    "        seen = set([target])\n",
    "        queue.append((target, 0))\n",
    "        ans = []\n",
    "        while queue:\n",
    "            node, dist = queue.popleft()\n",
    "            if dist == K:\n",
    "                ans.append(node.val)\n",
    "            for child in graph[node]:\n",
    "                if child not in seen:\n",
    "                    queue.append((child, dist+1))\n",
    "                    seen.add(child)\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, x):\n",
    "#         self.val = x\n",
    "#         self.left = None\n",
    "#         self.right = None\n",
    "\n",
    "class Solution:\n",
    "    def distanceK(self, root: TreeNode, target: TreeNode, K: int) -> List[int]:\n",
    "        def buildpar(root):\n",
    "            if root.left:\n",
    "                root.left.par = root\n",
    "                buildpar(root.left)\n",
    "            if root.right:\n",
    "                root.right.par = root\n",
    "                buildpar(root.right)\n",
    "        buildpar(root)\n",
    "        root.par = None\n",
    "        ##BFS\n",
    "        queue = collections.deque()\n",
    "        queue.append((target,0))\n",
    "        seen = []\n",
    "        while queue:\n",
    "            if queue[0][1] == K:\n",
    "                return [i[0].val for i in queue]\n",
    "            rootset = queue.popleft()\n",
    "            root = rootset[0]\n",
    "            depth = rootset[1]\n",
    "\n",
    "            seen.append(root)\n",
    "            for node in [root.par,root.left,root.right]:\n",
    "                if node and node not in seen:\n",
    "                    queue.append((node,depth+1))\n",
    "        return []\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, x):\n",
    "#         self.val = x\n",
    "#         self.left = None\n",
    "#         self.right = None\n",
    "\n",
    "class Solution:\n",
    "    def distanceK(self, root: TreeNode, target: TreeNode, K: int) -> List[int]:\n",
    "        def dfs(node, par = None):\n",
    "            if node:\n",
    "                node.par = par\n",
    "                dfs(node.left, node)\n",
    "                dfs(node.right, node)\n",
    "\n",
    "        dfs(root)\n",
    "\n",
    "        queue = collections.deque([(target, 0)])\n",
    "        seen = {target}\n",
    "        while queue:\n",
    "            if queue[0][1] == K:\n",
    "                return [node.val for node, d in queue]\n",
    "            node, d = queue.popleft()\n",
    "            for nei in (node.left, node.right, node.par):\n",
    "                if nei and nei not in seen:\n",
    "                    seen.add(nei)\n",
    "                    queue.append((nei, d+1))\n",
    "\n",
    "        return []\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from typing import List\n",
    "from copy import deepcopy\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "class TreeNode:\n",
    "    def __init__(self, x):\n",
    "        self.val = x\n",
    "        self.left = None\n",
    "        self.right = None\n",
    "\n",
    "class Solution:\n",
    "    def distanceK(self, root: TreeNode, target: TreeNode, K: int) -> List[int]:\n",
    "        def get_path(root, target):\n",
    "            nonlocal final_record\n",
    "            if not root:\n",
    "                return\n",
    "            if root.val != target.val:\n",
    "                record.append([root, 1])\n",
    "                get_path(root.left, target)\n",
    "                record.pop()\n",
    "                record.append([root, 2])\n",
    "                get_path(root.right, target)\n",
    "                record.pop()\n",
    "            else:\n",
    "                final_record = deepcopy(record)\n",
    "                return\n",
    "\n",
    "        def get_child_node(nodes, dis, k):\n",
    "            ans = []\n",
    "            if not nodes:\n",
    "                return nodes\n",
    "            if dis < k:\n",
    "                for node in nodes:\n",
    "                    if node.left:\n",
    "                        ans.append(node.left)\n",
    "                    if node.right:\n",
    "                        ans.append(node.right)\n",
    "                return get_child_node(ans, dis+1, k)\n",
    "            else:\n",
    "                return nodes\n",
    "\n",
    "\n",
    "        record = []\n",
    "        final_record = []\n",
    "        get_path(root, target)\n",
    "        # print(final_record)\n",
    "\n",
    "        ks1 = [target]\n",
    "        ks1 = get_child_node(ks1, 0, K)\n",
    "        ks1_val = [_.val for _ in ks1]\n",
    "\n",
    "        if len(final_record) >= K and K != 0:\n",
    "            ks1_val.append(final_record[-K][0].val)\n",
    "        # elif K == 0:\n",
    "        #     ks1_val.append(target.val)\n",
    "        # else:\n",
    "        #     pass\n",
    "        # print(final_record)\n",
    "        # print(final_record[::-1])\n",
    "        head_dis = 0\n",
    "        for head_record in final_record[::-1]:\n",
    "            # print(head_record)\n",
    "            head_dis += 1\n",
    "            remain_dis = K - head_dis - 1\n",
    "            # print(head_record[0].val, remain_dis)\n",
    "            if remain_dis >= 0:\n",
    "                head_node, direction = head_record\n",
    "                if direction == 1 and head_node.right:\n",
    "                    new_node = head_node.right\n",
    "                    to_add_nodes = get_child_node([new_node], 0, remain_dis)\n",
    "                    for node in to_add_nodes:\n",
    "                        ks1_val.append(node.val)\n",
    "                elif direction == 2 and head_node.left:\n",
    "                    new_node = head_node.left\n",
    "                    to_add_nodes = get_child_node([new_node], 0, remain_dis)\n",
    "                    for node in to_add_nodes:\n",
    "                        ks1_val.append(node.val)\n",
    "                else:\n",
    "                    pass\n",
    "\n",
    "        return ks1_val\n",
    "\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    pass\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, x):\n",
    "#         self.val = x\n",
    "#         self.left = None\n",
    "#         self.right = None\n",
    "\n",
    "class Solution:\n",
    "    def distanceK(self, root, target, K):\n",
    "        def dfs(node, par = None):\n",
    "            if node:\n",
    "                node.par = par\n",
    "                node.is_visit = False\n",
    "                node.distance = -1\n",
    "                if target.val == node.val:\n",
    "                    node.distance = 0\n",
    "                    node.is_visit = True\n",
    "                dfs(node.left, node)\n",
    "                dfs(node.right, node)\n",
    "\n",
    "        dfs(root)\n",
    "\n",
    "        queue = collections.deque([target])\n",
    "        # seen = {target}\n",
    "        d  = target.distance\n",
    "        while queue:\n",
    "            # if queue[0][1] == K:\n",
    "            #     return [node.val for node in queue]\n",
    "            if queue[0].distance == K:\n",
    "                return [node.val for node in queue]\n",
    "            node = queue.popleft()\n",
    "            d = node.distance + 1\n",
    "            for nei in (node.left, node.right, node.par):\n",
    "                if nei and nei.is_visit is False:\n",
    "                    nei.is_visit = True\n",
    "                    nei.distance = d\n",
    "                    queue.append(nei)\n",
    "\n",
    "        return []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, x):\n",
    "#         self.val = x\n",
    "#         self.left = None\n",
    "#         self.right = None\n",
    "\n",
    "class Solution:\n",
    "    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:\n",
    "        def parent_mem(root, parent_dict):\n",
    "            if root is None:\n",
    "                return\n",
    "            if root.left:\n",
    "                parent_dict[root.left.val] = root\n",
    "                parent_mem(root.left, parent_dict)\n",
    "            if root.right:\n",
    "                parent_dict[root.right.val] = root\n",
    "                parent_mem(root.right, parent_dict)\n",
    "            return \n",
    "        parent_dict = {}\n",
    "        parent_mem(root, parent_dict)\n",
    "        def find_k(target, from_node, distance, k, res, parent_dict):\n",
    "            if target is None:\n",
    "                return\n",
    "            if distance == k:\n",
    "                res.append(target.val)\n",
    "                return\n",
    "            if target.left != from_node:\n",
    "                find_k(target.left, target, distance + 1, k, res, parent_dict)\n",
    "            if target.right != from_node:\n",
    "                find_k(target.right, target, distance + 1, k, res, parent_dict)\n",
    "            if target.val in parent_dict and parent_dict[target.val] != from_node:\n",
    "                find_k(parent_dict[target.val], target, distance + 1, k, res, parent_dict)\n",
    "            return \n",
    "        res = []\n",
    "        find_k(target, None, 0, k, res, parent_dict)\n",
    "        return res\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, x):\n",
    "#         self.val = x\n",
    "#         self.left = None\n",
    "#         self.right = None\n",
    "\n",
    "from collections import defaultdict\n",
    "\n",
    "class TreeNode:\n",
    "    def __init__(self, x):\n",
    "        self.val = x\n",
    "        self.left = None\n",
    "        self.right = None\n",
    "\n",
    "def tree_to_graph(root):\n",
    "    graph = defaultdict(list)\n",
    "    \n",
    "    def traverse(node):\n",
    "        if node:\n",
    "            if node.left:\n",
    "                graph[node.val].append(node.left.val)\n",
    "                graph[node.left.val].append(node.val)\n",
    "                traverse(node.left)\n",
    "            if node.right:\n",
    "                graph[node.val].append(node.right.val)\n",
    "                graph[node.right.val].append(node.val)\n",
    "                traverse(node.right)\n",
    "    \n",
    "    traverse(root)\n",
    "    return graph\n",
    "\n",
    "class Solution:\n",
    "    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:\n",
    "        g = tree_to_graph(root)\n",
    "        vst = set()\n",
    "        cur = [target.val]\n",
    "        vst.add(target.val)\n",
    "        for _ in range(k):\n",
    "            nxt = []\n",
    "            for v in cur:\n",
    "                for w in g[v]:\n",
    "                    if w not in vst:\n",
    "                        vst.add(w)\n",
    "                        nxt.append(w)\n",
    "            cur = nxt\n",
    "        return cur"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, x):\n",
    "#         self.val = x\n",
    "#         self.left = None\n",
    "#         self.right = None\n",
    "\n",
    "class Solution:\n",
    "    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:\n",
    "        from collections import defaultdict\n",
    "        graph = defaultdict(set)\n",
    "        # 建图\n",
    "        def dfs(root):\n",
    "            if root.left :\n",
    "                graph[root.val].add(root.left.val)\n",
    "                graph[root.left.val].add(root.val)\n",
    "                dfs(root.left)\n",
    "            if root.right:\n",
    "                graph[root.val].add(root.right.val)\n",
    "                graph[root.right.val].add(root.val)\n",
    "                dfs(root.right)\n",
    "        dfs(root)\n",
    "        #print(graph)\n",
    "        cur = [target.val]\n",
    "        visited ={target.val}\n",
    "        while k:\n",
    "            next_time = []\n",
    "            while cur:\n",
    "                tmp = cur.pop()\n",
    "                for node in graph[tmp]:\n",
    "                    if node not in visited:\n",
    "                        visited.add(node)\n",
    "                        next_time.append(node)\n",
    "            k -= 1\n",
    "            cur = next_time\n",
    "        return cur\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, x):\n",
    "#         self.val = x\n",
    "#         self.left = None\n",
    "#         self.right = None\n",
    "\n",
    "class Solution:\n",
    "    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:\n",
    "        mp = defaultdict(list)\n",
    "        \n",
    "        # DFS 树转图\n",
    "        st = [root]\n",
    "        while st:\n",
    "            node = st.pop()\n",
    "            if node.left:\n",
    "                mp[node.val].append(node.left.val)\n",
    "                mp[node.left.val].append(node.val)\n",
    "                st.append(node.left)\n",
    "            if node.right:\n",
    "                mp[node.val].append(node.right.val)\n",
    "                mp[node.right.val].append(node.val)\n",
    "                st.append(node.right)\n",
    "        print(mp)\n",
    "        # BFS 搜索target 层数为2\n",
    "        q = [target.val]\n",
    "        vis = [False] * 1000\n",
    "        vis[target.val] = True\n",
    "        step = 0\n",
    "        while q:\n",
    "            tmp = q\n",
    "            q = []\n",
    "            if step == k:\n",
    "                return tmp\n",
    "            for x in tmp:\n",
    "                for y in mp[x]:\n",
    "                    if not vis[y]:\n",
    "                        q.append(y)\n",
    "                        vis[y] = True\n",
    "            step += 1\n",
    "        return []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, x):\n",
    "#         self.val = x\n",
    "#         self.left = None\n",
    "#         self.right = None\n",
    "\n",
    "class Solution:\n",
    "    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:\n",
    "        if k == 0:\n",
    "            return [target.val]\n",
    "        graph = [ [] for _ in range(501)]\n",
    "        def dfs(u:int):\n",
    "            if u.left != None:\n",
    "                graph[u.val].append(u.left.val)\n",
    "                graph[u.left.val].append(u.val)\n",
    "                dfs(u.left)\n",
    "            if u.right != None:\n",
    "                graph[u.val].append(u.right.val)\n",
    "                graph[u.right.val].append(u.val)\n",
    "                dfs(u.right)\n",
    "        dfs(root)\n",
    "        # print(graph)\n",
    "        q = [target.val]\n",
    "        dis = [inf] * 501\n",
    "        dis[target.val] = 0\n",
    "        vis = [False] * 501\n",
    "        vis[target.val] = True\n",
    "        ans = []\n",
    "        while q:\n",
    "            nxts = []\n",
    "            for u in q:\n",
    "                for v in graph[u]:\n",
    "                    if not vis[v]:\n",
    "                        dis[v] = dis[u] + 1\n",
    "                        if dis[v] == k:\n",
    "                            ans.append(v)\n",
    "                        else:\n",
    "                            vis[v] = True\n",
    "                            nxts.append(v)\n",
    "            q = nxts\n",
    "        return ans\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, x):\n",
    "#         self.val = x\n",
    "#         self.left = None\n",
    "#         self.right = None\n",
    "\n",
    "\n",
    "import copy\n",
    "\n",
    "def findAllPaths(root, pathDict, target, k, level, flag, retArr, parentNode, rootNode, targetLevel):\n",
    "    if not root:\n",
    "        return\n",
    "    elif root.left == None and root.right == None:\n",
    "        if root.val not in pathDict:\n",
    "            pathDict[root.val] = [level, [[root.val]]]\n",
    "\n",
    "        if flag[0] == 1 and level - pathDict[target.val][0] == k:\n",
    "            retArr.append(root.val)\n",
    "        elif flag[0] == 4 and level == targetLevel[0]:\n",
    "            retArr.append(root.val)\n",
    "        elif flag[0] == 4 and level > targetLevel[0]:\n",
    "            return    \n",
    "\n",
    "        if root.val == target.val:\n",
    "            if k:\n",
    "                flag[0] = 2\n",
    "            else:\n",
    "                retArr.append(root.val)\n",
    "                flag[0] = 3\n",
    "\n",
    "        return\n",
    "\n",
    "    if root.val not in pathDict:\n",
    "        pathDict[root.val] = [level, []]\n",
    "\n",
    "    if root.val == target.val:\n",
    "        if k:\n",
    "            flag[0] = 1\n",
    "        else:\n",
    "            retArr.append(root.val)\n",
    "            flag[0] = 3\n",
    "            return\n",
    "\n",
    "    if flag[0] == 1:\n",
    "        if level - pathDict[target.val][0] == k:\n",
    "            retArr.append(root.val)\n",
    "            return\n",
    "    elif flag[0] == 4 and level == targetLevel[0]:\n",
    "        retArr.append(root.val)\n",
    "    elif flag[0] == 4 and level > targetLevel[0]:\n",
    "        return\n",
    "\n",
    "    if root.left:\n",
    "        findAllPaths(root.left, pathDict, target, k, level+1, flag, retArr, parentNode, rootNode, targetLevel)\n",
    "        if flag[0] == 3:\n",
    "            return\n",
    "        if flag[0] == 0:\n",
    "            for x in pathDict[root.left.val][1]:\n",
    "                tmp = copy.deepcopy(x)\n",
    "                tmp.append(root.val)\n",
    "                pathDict[root.val][1].append(tmp)\n",
    "\n",
    "        if flag[0] == 2:\n",
    "            if pathDict[target.val][0] - level == k:\n",
    "                retArr.append(root.val)\n",
    "                flag[0] = 3\n",
    "                return\n",
    "            else:\n",
    "                count = 0\n",
    "                for x in pathDict[root.val][1]:\n",
    "                    if root.left and root.left.val not in x:\n",
    "                        if len(x) >= k + level - pathDict[target.val][0] + 1:\n",
    "                            idx = k + level - pathDict[target.val][0]\n",
    "                            if x[-idx-1] not in retArr:\n",
    "                                retArr.append(x[-idx-1])\n",
    "                            count += 1\n",
    "\n",
    "                if count == 0:\n",
    "                    targetLevel[0] = 2 * level + k - pathDict[target.val][0]\n",
    "                    parentNode[0] = root.val\n",
    "                    flag[0] = 4\n",
    "\n",
    "            # if root == rootNode:\n",
    "            #     flag[0] = 3\n",
    "            #     return\n",
    "\n",
    "\n",
    "    if root.right:\n",
    "        findAllPaths(root.right, pathDict, target, k, level+1, flag, retArr, parentNode, rootNode, targetLevel)\n",
    "        if flag[0] == 3:\n",
    "            return\n",
    "        if flag[0] == 0:\n",
    "            for x in pathDict[root.right.val][1]:\n",
    "                tmp = copy.deepcopy(x)\n",
    "                tmp.append(root.val)\n",
    "                pathDict[root.val][1].append(tmp)\n",
    "\n",
    "\n",
    "\n",
    "    if root.val == target.val:\n",
    "        flag[0] = 2\n",
    "        parentNode[0] = root.val\n",
    "    elif flag[0] == 2:\n",
    "        if pathDict[target.val][0] - level == k:\n",
    "            retArr.append(root.val)\n",
    "            flag[0] = 3\n",
    "            return\n",
    "        else:\n",
    "            count = 0\n",
    "            for x in pathDict[root.val][1]:\n",
    "                if root.right and root.right.val not in x:\n",
    "                    if len(x) >= k + level - pathDict[target.val][0] + 1:\n",
    "                        idx = k + level - pathDict[target.val][0]\n",
    "                        if x[-idx-1] not in retArr:\n",
    "                            retArr.append(x[-idx-1])\n",
    "                        count += 1\n",
    "\n",
    "            # for x in pathDict[root.val][1]:\n",
    "            #     if len(x) >= k + level - pathDict[target][0]:\n",
    "            #         idx = k + level - pathDict[target][0]\n",
    "            #         retArr.append(x[-idx])\n",
    "\n",
    "    elif flag[0] == 4:\n",
    "        if parentNode[0] == root.val:\n",
    "            flag[0] = 2\n",
    "\n",
    "class Solution(object):\n",
    "    def distanceK(self, root, target, k):\n",
    "        \"\"\"\n",
    "        :type root: TreeNode\n",
    "        :type target: TreeNode\n",
    "        :type k: int\n",
    "        :rtype: List[int]\n",
    "        \"\"\"\n",
    "\n",
    "        pathDict = {}\n",
    "        retArr = []\n",
    "        flag = [0]\n",
    "        parentNode = [-1]\n",
    "        targetLevel = [-1]\n",
    "\n",
    "        findAllPaths(root, pathDict, target, k, 1, flag, retArr, parentNode, root, targetLevel)\n",
    "        # print(pathDict)\n",
    "        # print(retArr)\n",
    "        # graphDict = {}\n",
    "        # preOrderTraverse(root, graphDict)\n",
    "\n",
    "        # print(graphDict)\n",
    "        #\n",
    "        # return BFS(graphDict, k, target)\n",
    "        return retArr\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, x):\n",
    "#         self.val = x\n",
    "#         self.left = None\n",
    "#         self.right = None\n",
    "\n",
    "class Solution:\n",
    "    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:\n",
    "        self.k = k\n",
    "        self.target = target\n",
    "        self.vals = []\n",
    "        self.find(root)\n",
    "        return self.vals\n",
    "\n",
    "    def find(self, root, ins=0):\n",
    "        if not root:\n",
    "            return 0\n",
    "        elif root.val == self.target.val:\n",
    "            if self.k == 0:\n",
    "                self.vals.append(self.target.val)\n",
    "                return 0\n",
    "            self.find(root.left, ins + 1)\n",
    "            self.find(root.right, ins + 1)\n",
    "            return 1\n",
    "        elif self.k and ins == self.k:\n",
    "            self.vals.append(root.val)\n",
    "            return 0\n",
    "        elif ins:\n",
    "            self.find(root.left, ins + 1)\n",
    "            self.find(root.right, ins + 1)\n",
    "            return 0\n",
    "        else:\n",
    "            outs = self.find(root.left, ins)\n",
    "            if outs:\n",
    "                if outs == self.k:\n",
    "                    self.vals.append(root.val)\n",
    "                    return 0\n",
    "                self.find(root.right, outs+1)\n",
    "                return outs + 1\n",
    "            outs = self.find(root.right, ins)\n",
    "            if outs:\n",
    "                if outs + ins == self.k:\n",
    "                    self.vals.append(root.val)\n",
    "                    return 0\n",
    "                self.find(root.left, outs+1)\n",
    "                return outs + 1\n",
    "            return 0\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, x):\n",
    "#         self.val = x\n",
    "#         self.left = None\n",
    "#         self.right = None\n",
    "\n",
    "class Solution:\n",
    "    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:\n",
    "        def find_parent(root, parent):\n",
    "            if not root: return\n",
    "            parent_dict[root] = parent\n",
    "            find_parent(root.left, root)\n",
    "            find_parent(root.right, root)\n",
    "        \n",
    "        def dfs(root, pre, depth):\n",
    "            if not root: return\n",
    "            if depth == k:\n",
    "                res.append(root.val)\n",
    "            if root.left and root.left != pre:\n",
    "                dfs(root.left, root, depth+1)\n",
    "            if root.right and root.right != pre:\n",
    "                dfs(root.right, root, depth+1)\n",
    "            if parent_dict[root] and parent_dict[root] != pre:\n",
    "                dfs(parent_dict[root], root, depth+1)\n",
    "\n",
    "        parent_dict = {}\n",
    "        find_parent(root, None)\n",
    "        res = []\n",
    "        dfs(target, None, 0)\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from collections import defaultdict\n",
    "class TreeNode:\n",
    "    def __init__(self, x):\n",
    "        self.val = x\n",
    "        self.left = None\n",
    "        self.right = None\n",
    "\n",
    "def listToTree(ls):\n",
    "    length = len(ls)\n",
    "    if length <= 0:\n",
    "        return None\n",
    "    root = TreeNode(ls[0])\n",
    "    queue = [root]\n",
    "\n",
    "    listIndex = 1\n",
    "    while(listIndex < length):\n",
    "        if not queue:\n",
    "            break\n",
    "        node = queue.pop(0)\n",
    "        # 左子树\n",
    "        if ls[listIndex] != \"null\":\n",
    "            node.left = TreeNode(ls[listIndex])\n",
    "            queue.append(node.left)\n",
    "        listIndex += 1\n",
    "\n",
    "        if listIndex >= length:\n",
    "            break\n",
    "\n",
    "        # 右子树\n",
    "        if ls[listIndex] != \"null\":\n",
    "            node.right = TreeNode(ls[listIndex])\n",
    "            queue.append(node.right)\n",
    "        listIndex += 1\n",
    "\n",
    "    return root\n",
    "\n",
    "def display(root):\n",
    "    if root:\n",
    "        print(root.val)\n",
    "        display(root.left)\n",
    "        display(root.right)\n",
    "\n",
    "class Solution:\n",
    "    def distanceK(self, root, target, k):\n",
    "\n",
    "        graph = defaultdict(set)\n",
    "        self.dfs(root, graph)\n",
    "        # print(\"graph:\", graph)\n",
    "\n",
    "        # bfs遍历到某个节点距离K的节点\n",
    "        cur = [target.val]\n",
    "        visited = {target.val}\n",
    "        for i in range(k):\n",
    "            nextList = []\n",
    "            while cur:\n",
    "                tmp = cur.pop(0)\n",
    "                for node in graph[tmp]:\n",
    "                    if node in visited:\n",
    "                        continue\n",
    "                    visited.add(node)\n",
    "                    nextList.append(node)\n",
    "            cur = nextList\n",
    "        return cur\n",
    "\n",
    "    # 建图\n",
    "    def dfs(self, root, graph): # 树只有一到两个孩子，是稀疏图，所以用邻接表存储\n",
    "        if root.left:\n",
    "            graph[root.val].add(root.left.val)\n",
    "            graph[root.left.val].add(root.val)\n",
    "            self.dfs(root.left, graph)\n",
    "        if root.right:\n",
    "            graph[root.val].add(root.right.val)\n",
    "            graph[root.right.val].add(root.val)\n",
    "            self.dfs(root.right, graph)\n",
    "\n",
    "# if __name__ == \"__main__\":\n",
    "#     root = [3, 5, 1, 6, 2, 0, 8, \"null\", \"null\", 7, 4]\n",
    "#     target = TreeNode(5)\n",
    "#     k = 2\n",
    "#     root = listToTree(root)\n",
    "#     display(root)\n",
    "#     so = Solution()\n",
    "#     ans = so.distanceK(root, target, k)\n",
    "#     print(\"ans:\", ans)\n"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}
